"""
    ErrorInflation

Adaptive observation error inflation based on innovation statistics.
Ensures observation errors are well-calibrated for optimal analysis.

# Mathematical Foundation

## Chi-Square Statistic
```
χ² = Σᵢ (d_i / σ_i)²
```

Expected value: E[χ²] = N (number of observations)

## Normalized Chi-Square
```
χ²_norm = χ² / N
```

## Error Calibration

If χ²_norm ≈ 1: errors well-calibrated
If χ²_norm < 1: errors overestimated (observations too uncertain)
If χ²_norm > 1: errors underestimated (observations too confident)

## Inflation Factor
```
α = √(χ²_norm) = √(χ² / N)
```

Inflated error:
```
σ_inflated = α * σ_original
```

# Adaptive Approach

1. Compute χ² for observation subset (by type, region, etc.)
2. Calculate inflation factor α
3. Apply inflation to observation errors
4. Iterate until χ²_norm ≈ 1

# Applications

**Per observation type**: Different inflation for radiosonde, aircraft, satellite
**Per region**: Different inflation for tropics, mid-latitudes, poles
**Per variable**: Different inflation for T, u, v, q
**Temporal**: Adjust inflation during assimilation cycle
"""
module ErrorInflation

using LinearAlgebra
using Statistics
using ...ObservationTypes

# Export types and functions
export ErrorInflationConfig, ErrorInflationResult
export perform_error_inflation, compute_chi_square, compute_inflation_factor
export inflate_errors_adaptive

"""
    ErrorInflationConfig{T<:AbstractFloat}

Configuration for adaptive error inflation.
"""
struct ErrorInflationConfig{T<:AbstractFloat}
    target_chi_square::T          # Target χ²/N (typically 1.0)
    chi_square_tolerance::T       # Acceptable deviation from target
    max_inflation_factor::T       # Maximum allowed inflation (safety)
    min_inflation_factor::T       # Minimum allowed inflation
    min_obs_for_inflation::Int    # Minimum obs needed to compute inflation
    inflation_by_type::Bool       # Separate inflation per obs type
    inflation_by_region::Bool     # Separate inflation per region
    damping_factor::T            # Damp inflation changes (0-1)

    function ErrorInflationConfig{T}(;
        target_chi_square = T(1.0),
        chi_square_tolerance = T(0.1),
        max_inflation_factor = T(3.0),
        min_inflation_factor = T(0.5),
        min_obs_for_inflation = 50,
        inflation_by_type = true,
        inflation_by_region = false,
        damping_factor = T(0.7)
    ) where T
        new{T}(target_chi_square, chi_square_tolerance, max_inflation_factor,
               min_inflation_factor, min_obs_for_inflation, inflation_by_type,
               inflation_by_region, damping_factor)
    end
end

ErrorInflationConfig() = ErrorInflationConfig{Float64}()

"""
    ErrorInflationResult{T<:AbstractFloat}

Result of error inflation procedure.
"""
struct ErrorInflationResult{T<:AbstractFloat}
    inflated_observations::Vector{AbstractObservation{T}}
    inflation_factors::Dict{String,T}
    chi_square_before::T
    chi_square_after::T
    chi_square_per_obs_before::T
    chi_square_per_obs_after::T
    n_observations::Int
    converged::Bool
    statistics::Dict{String,Any}
end

"""
    perform_error_inflation(observations::Vector{<:AbstractObservation{T}},
                           background_values::Vector{T},
                           config::ErrorInflationConfig{T}) where T

Perform adaptive error inflation.
"""
function perform_error_inflation(observations::Vector{<:AbstractObservation{T}},
                                background_values::Vector{T},
                                config::ErrorInflationConfig{T}) where T
    n_obs = length(observations)

    if n_obs < config.min_obs_for_inflation || length(background_values) != n_obs
        # Not enough observations - return unchanged
        return ErrorInflationResult{T}(
            copy(observations),
            Dict{String,T}(),
            T(0), T(0), T(0), T(0),
            n_obs, false,
            Dict{String,Any}("reason" => "insufficient_observations")
        )
    end

    # Compute initial chi-square
    chi_sq_before = compute_chi_square(observations, background_values)
    chi_sq_per_obs_before = chi_sq_before / n_obs

    # Determine if inflation needed
    if abs(chi_sq_per_obs_before - config.target_chi_square) <= config.chi_square_tolerance
        # Already well-calibrated
        return ErrorInflationResult{T}(
            copy(observations),
            Dict{String,T}("overall" => T(1.0)),
            chi_sq_before, chi_sq_before,
            chi_sq_per_obs_before, chi_sq_per_obs_before,
            n_obs, true,
            Dict{String,Any}("reason" => "already_calibrated")
        )
    end

    # Compute inflation factors
    inflation_factors = if config.inflation_by_type
        compute_inflation_by_type(observations, background_values, config)
    else
        Dict{String,T}("overall" => compute_inflation_factor(chi_sq_per_obs_before, config))
    end

    # Apply inflation
    inflated_obs = inflate_errors(observations, inflation_factors, config)

    # Compute chi-square after inflation
    chi_sq_after = compute_chi_square(inflated_obs, background_values)
    chi_sq_per_obs_after = chi_sq_after / n_obs

    # Check convergence
    converged = abs(chi_sq_per_obs_after - config.target_chi_square) <= config.chi_square_tolerance

    stats = Dict{String,Any}(
        "chi_square_reduction" => chi_sq_before - chi_sq_after,
        "inflation_factors" => inflation_factors,
        "mean_inflation" => mean(values(inflation_factors)),
        "converged" => converged
    )

    return ErrorInflationResult{T}(
        inflated_obs,
        inflation_factors,
        chi_sq_before,
        chi_sq_after,
        chi_sq_per_obs_before,
        chi_sq_per_obs_after,
        n_obs,
        converged,
        stats
    )
end

"""
    compute_chi_square(observations, background_values)

Compute chi-square statistic for observations.
"""
function compute_chi_square(observations::Vector{<:AbstractObservation{T}},
                           background_values::Vector{T}) where T
    chi_sq = T(0)

    for (i, obs) in enumerate(observations)
        innovation = obs.observed_value - background_values[i]
        error = obs.observation_error.total_error

        normalized_innov = innovation / error
        chi_sq += normalized_innov^2
    end

    return chi_sq
end

"""
    compute_inflation_factor(chi_square_per_obs, config)

Compute inflation factor from chi-square statistic.
"""
function compute_inflation_factor(chi_sq_per_obs::T, config::ErrorInflationConfig{T}) where T
    # Inflation factor is sqrt of normalized chi-square
    factor = sqrt(chi_sq_per_obs / config.target_chi_square)

    # Apply damping to avoid oscillations
    factor = 1.0 + config.damping_factor * (factor - 1.0)

    # Clamp to reasonable range
    factor = clamp(factor, config.min_inflation_factor, config.max_inflation_factor)

    return factor
end

"""
    compute_inflation_by_type(observations, background_values, config)

Compute separate inflation factors for each observation type.
"""
function compute_inflation_by_type(observations::Vector{<:AbstractObservation{T}},
                                  background_values::Vector{T},
                                  config::ErrorInflationConfig{T}) where T
    # Group observations by type
    type_groups = Dict{String,Vector{Int}}()

    for (i, obs) in enumerate(observations)
        obs_type = get_observation_type_name(obs)

        if haskey(type_groups, obs_type)
            push!(type_groups[obs_type], i)
        else
            type_groups[obs_type] = [i]
        end
    end

    # Compute inflation for each type
    inflation_factors = Dict{String,T}()

    for (obs_type, indices) in type_groups
        if length(indices) < config.min_obs_for_inflation
            # Not enough obs of this type - use overall inflation
            inflation_factors[obs_type] = T(1.0)
            continue
        end

        # Compute chi-square for this type
        chi_sq = T(0)
        for idx in indices
            obs = observations[idx]
            innovation = obs.observed_value - background_values[idx]
            error = obs.observation_error.total_error
            chi_sq += (innovation / error)^2
        end

        chi_sq_per_obs = chi_sq / length(indices)

        # Compute inflation factor
        factor = compute_inflation_factor(chi_sq_per_obs, config)
        inflation_factors[obs_type] = factor
    end

    return inflation_factors
end

"""
    inflate_errors(observations, inflation_factors, config)

Apply inflation factors to observation errors.
"""
function inflate_errors(observations::Vector{<:AbstractObservation{T}},
                       inflation_factors::Dict{String,T},
                       config::ErrorInflationConfig{T}) where T
    inflated_obs = similar(observations)

    for (i, obs) in enumerate(observations)
        obs_type = get_observation_type_name(obs)

        # Get inflation factor for this type
        factor = get(inflation_factors, obs_type, get(inflation_factors, "overall", T(1.0)))

        # Create inflated observation error
        old_error = obs.observation_error
        new_total_error = old_error.total_error * factor

        inflated_error = ObservationError{T}(
            old_error.instrument_error * factor,
            old_error.representativeness_error * factor,
            old_error.preprocessing_error * factor,
            old_error.error_correlation_length,
            old_error.error_correlation_time
        )

        # Create new observation with inflated error
        # (This is simplified - would need proper constructor for each type)
        inflated_obs[i] = obs  # Placeholder - in practice would reconstruct observation
    end

    return inflated_obs
end

"""
    get_observation_type_name(obs)

Get observation type name for grouping.
"""
function get_observation_type_name(obs::SurfaceObservation)::String
    return "surface_$(obs.variable)"
end

function get_observation_type_name(obs::RadiosondeObservation)::String
    return "radiosonde_$(obs.variable)"
end

function get_observation_type_name(obs::AircraftObservation)::String
    return "aircraft_$(obs.variable)"
end

function get_observation_type_name(obs::RadianceObservation)::String
    return "radiance_$(obs.satellite_id)_ch$(obs.channel)"
end

function get_observation_type_name(obs::GPSROObservation)::String
    return "gpsro_$(obs.variable)"
end

function get_observation_type_name(obs::RadarObservation)::String
    return "radar_$(obs.variable)"
end

function get_observation_type_name(obs::AbstractObservation)::String
    return "unknown_$(obs.variable)"
end

"""
    inflate_errors_adaptive(observations, background_values, max_iterations=5)

Iteratively inflate errors until chi-square converges.
"""
function inflate_errors_adaptive(observations::Vector{<:AbstractObservation{T}},
                                background_values::Vector{T},
                                max_iterations::Int=5) where T
    config = ErrorInflationConfig{T}()
    current_obs = copy(observations)

    for iter in 1:max_iterations
        result = perform_error_inflation(current_obs, background_values, config)

        if result.converged
            return result
        end

        current_obs = result.inflated_observations
    end

    # Return last iteration result
    return perform_error_inflation(current_obs, background_values, config)
end

end # module ErrorInflation
